import torch
import torch.nn as nn 


class Conv(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, padding=1) -> None:
        super().__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, 1, 1, bias=True),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(),
        )
    def forward(self, x):
        return self.conv(x)

class TwoConv(nn.Module):
    def __init__(self, in_channels, out_channels) -> None:
        super().__init__()
        self.conv1 = Conv(in_channels, out_channels)
        self.conv2 = Conv(out_channels, out_channels)

    def forward(self, x):
        return self.conv2(self.conv1(x))

 
class Upsample(nn.Module):
    def __init__(self, in_channels, out_channels) -> None:
        super().__init__()
        self.ups = nn.Upsample(scale_factor=2, mode="trilinear")
        self.conv_up = TwoConv(in_channels, out_channels)
    
    def forward(self, x, x_concat):
        x = self.ups(x)
        x = torch.cat([x, x_concat], dim=1)

        x = self.conv_up(x)
        